#ifndef SKIPGRID_H
#define SKIPGRID_H

/**
 * 参考: GREGORIO D, STEFANO L. SkiMap: An Efficient Mapping Framework for Robot Navigation[J]. arXiv: Computer Vision and Pattern Recognition, 2017.
 */

#include "skiplist.h"
#include <omp.h>
#include <functional>

// 修改1

/*
 * 嵌套二维表结构，已知xy索引时返回SkipNode指针，未知索引时返回GridNode
 */
template <typename KEY, typename VAL>
struct GridNode
{
    KEY ix; // x向索引
    KEY iy; // y向索引
    VAL val; 

    GridNode(KEY ix_, KEY iy_, const VAL& val_) : ix(ix_), iy(iy_), val(val_) {}
};

template <typename KEY, typename VAL, unsigned LEVEL = 10>
class SkipGrid
{
public:
    using YBranch = SkipList<KEY, VAL, LEVEL>; // Y分支类型
    using XBranch = SkipList<KEY, YBranch*, LEVEL>; //X分支类型
    using Node = SkipNode<KEY, VAL, LEVEL>; 
    using Grid = struct GridNode<KEY, VAL>;

    SkipGrid(KEY min_key, KEY max_key)
    {
        // 初始化根节点
        _root = new XBranch(min_key, max_key);
    }

    ~SkipGrid()
    {
        // 删除Y分支
        auto x_node = _root->head()->forward[1];
        
        while (x_node != _root->tail())
        {
            if (x_node->val)
            {
                delete x_node->val;
            }
            x_node = x_node->forward[1];
        }

        // 删除X分支
        delete _root;
    }

    // 搜索节点
    const Node* search(KEY ix, KEY iy) const 
    {
        // 搜索X分支
        auto x_node = _root->search(ix);

        // 搜索Y分支
        if (x_node && x_node->val)
        {
            auto y_node = x_node->val->search(iy);
            return y_node;
        }

        return nullptr;
    }

    // 插入&更新节点
    Node* integrate(KEY ix, KEY iy, const VAL& val)
    {
        // 搜索X分支
        auto x_node = _root->search(ix);

        // 若分支不存在则创建
        if (!x_node)
            x_node = _root->integrate(ix, new YBranch(_root->head()->key, _root->tail()->key));

        // 创建失败返回空
        if (!x_node)
            return nullptr;
    
        // 更新Y分支
        return x_node->val->integrate(iy, val);
    }

    // 移除元素
    void remove(KEY ix, KEY iy)
    {
        // 搜索X分支
        auto x_node = _root->search(ix);

        // 若分支存在，尝试删除
        if (x_node && x_node->val)
        {
            x_node->val->remove(iy);

            // 若Y分支无元素，则删除整个分支
            if (x_node->val->size() == 0)
            {
                delete x_node->val;
                _root->remove(ix);
            }
        }
    }

    // 二维表 --> 向量
    void toVector(std::vector<Grid>& vec)
    {
        vec.clear();

        // 获取X分支列表
        std::vector<typename XBranch::Node*> x_nodes;
        _root->toVector(x_nodes);

        // 并行遍历Y分支
        #pragma omp parallel
        {
            std::vector<Grid> grids;

            #pragma omp for nowait
            for (unsigned i = 0; i < x_nodes.size(); ++i)
            {
                std::vector<Node*> y_nodes;
                if (x_nodes[i] && x_nodes[i]->val)
                {
                    x_nodes[i]->val->toVector(y_nodes);
                    for (unsigned j = 0; j < y_nodes.size(); ++j)
                    {
                        grids.emplace_back(Grid{x_nodes[i]->key, y_nodes[j]->key, y_nodes[j]->val});
                    }
                }
            }

            // 临界区
            #pragma omp critical
            vec.insert(vec.end(), grids.begin(), grids.end());
        }
    }

    // 二维表 --> 向量，且满足：left_key <= ix <= right_key && bottom_key <= iy <= top_key && filter(ix, iy) == true
    void toVector(std::vector<Grid>& vec, KEY left_key, KEY right_key, KEY bottom_key, KEY top_key, std::function<bool(KEY,KEY)> filter = nullptr)
    {
        vec.clear();

        // 获取X分支列表
        std::vector<typename XBranch::Node*> x_nodes;
        _root->toVector(x_nodes, left_key, right_key);

        // 并行遍历Y分支
        #pragma omp parallel
        {
            std::vector<Grid> grids;

            #pragma omp for nowait
            for (unsigned i = 0; i < x_nodes.size(); ++i)
            {
                std::vector<Node*> y_nodes;
                if (x_nodes[i] && x_nodes[i]->val)
                {
                    x_nodes[i]->val->toVector(y_nodes, bottom_key, top_key);
                    for (unsigned j = 0; j < y_nodes.size(); ++j)
                    {
                        if (filter != nullptr)
                        {
                            if (filter(x_nodes[i]->key, y_nodes[j]->key))
                            {
                                grids.emplace_back(Grid{x_nodes[i]->key, y_nodes[j]->key, y_nodes[j]->val});
                            }
                        } 
                        else
                        {
                            grids.emplace_back(Grid{x_nodes[i]->key, y_nodes[j]->key, y_nodes[j]->val});
                        }
                    }
                }
            }

            // 临界区
            #pragma omp critical
            vec.insert(vec.end(), grids.begin(), grids.end());
        }
    }

private:
    // X分支根节点
    // XBranch* _root = nullptr;
    XBranch* _root; // 修改3
};

// 修改2

#endif